#!/usr/bin/env python
# coding: utf-8

# In[ ]:

import os
import numpy as np
import pandas as pd
from pathos.multiprocessing import Pool, cpu_count
import emcee
from functools import partial

from operator import add
from functools import reduce
from contextlib import closing

from jm.library.data_helper import filter_data, is_between
from jm.library.likelihood import Likelihood
from jm.library.simulator import TargetPrioritiesTargetActions, CounterfactualConfigs
from jm.library.additional_helpers import payment_cols, priority_cols, action_cols, \
    relative_payment_cols, fwdiff_payment_cols, relative_fwdiff_payment_cols, \
    cumul_binary_fwdiff_payment_cols, binary_fwdiff_payment_cols, get_cols_over_time


# In[ ]:
df_sim = pd.read_csv('estimation/simulation_estimates/longer_G1_more_G1_actual_estimates.csv')
df_sim_control = pd.read_csv('estimation/simulation_estimates/control_config_control_total_due_estimates.csv')
df_status = pd.read_csv('data/jesus_maria_status_deidentified.csv')

df_status[relative_payment_cols] = df_status[payment_cols].divide(
                df_status['total_due'], axis=0)


df_status[relative_fwdiff_payment_cols] = df_status[fwdiff_payment_cols].divide(
    df_status['total_due'], axis=0)

df_treatment = filter_data(df_status, assignment_to_treatment=1)
df_control = filter_data(df_status, assignment_to_treatment=0)

df_treatment = df_treatment.sort_values(
    'effective_score', ascending=False, kind='mergesort')
state_0_selected_cols = ['total_due'] + [payment_cols[0], priority_cols[0], action_cols[0]] + [
    'prob_repayment_endo_covariates']


td_pctile99 = df_status['total_due'].quantile([0.99]).values[0]


# ### Number of Payment Events in Control v.s. Treatment

with open('figs/tableOB1_sim_longversion.txt', 'w') as f:
    f.write('Base Num Event:')
    f.write(str(df_sim['cumul_binary_events'].mean()))
    f.write('\n')
    f.write('Control Num Event:')
    f.write(str(df_sim_control['cumul_binary_events'].mean()))
    f.write('\n')

with open('figs/tableOB2_sim_longversion.txt', 'w') as f:
    f.write('Base Average Payment Per Event:')
    f.write(str(df_sim['payment_per_event'].mean()))
    f.write('\n')
    f.write('Control Average Payment Per Event:')
    f.write(str(df_sim_control['payment_per_event'].mean()))
    f.write('\n')


# #### <99th percentile

with open('figs/tableOB2_sim_longversion.txt', 'a') as f:
    f.write('Base Average Payment Per Event (<99p):')
    f.write(str(df_sim['payment_per_event_p99'].mean()))
    f.write('\n')
    f.write('Control Average Payment Per Event (<99p):')
    f.write(str(df_sim_control['payment_per_event_p99'].mean()))
    f.write('\n')
